# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import logging
from typing import List, Union

from tornado import gen

from coin.base.timestamp import get_timestamp
import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.kraken.kr_rest.native_public_client import KrakenNativePublicClient
from coin.exchange.kraken.rest.async_feed_client import KrakenAsyncFeedClient


class RestJsonSubscriberAsOfWebsocket(WebsocketJsonSubscriberBase):
  pass


class KrakenSubscriber(RestJsonSubscriberAsOfWebsocket):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    logger = logger or logging.getLogger(__name__)
    super().__init__(close_callback, logger)
    self._symbols = to_list(symbol)
    assert len(self._symbols) >= 1, self._symbols

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

    self._async_feed_client = KrakenAsyncFeedClient(KrakenNativePublicClient(timeout=5), logger)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('kraken_rest_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('kraken_rest_trade_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'orderbook': queue.get_writer(topic_id=topic_depth[1]),
        'transaction': queue.get_writer(topic_id=topic_trade[1]),
    }

    self._enable_topic_event_publisher(queue,
                                       'kraken_rest_topic_event', [topic_depth[0], topic_trade[0]])
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_order_book(self, fut):
    if self.stopped:
      return
    update = fut.result()
    if update is not None:
      message = update.msg
      timestamp = get_timestamp()
      self._publisher['orderbook'].write(timestamp, message)

  def on_trade(self, fut):
    if self.stopped:
      return
    update = fut.result()
    if update is not None:
      message = update.msg
      timestamp = get_timestamp()
      self._publisher['transaction'].write(timestamp, message)

  @gen.coroutine
  def _run_loop(self):
    try:
      self.on_open()
      ns_before = 30 * 10**9
      while self._state == self.State.RUNNING:
        # Query order book every half second for each symbol.
        for _ in range(2):
          for symbol in self._symbols:
            # Query order book of 30 levels.
            fut = self._async_feed_client.async_query_order_book(symbol, level=30)
            if fut is not None:
              self._ioloop.add_future(fut, self.on_order_book)
          yield gen.sleep(0.5)
        # Query trade every one second for each symbol.
        for symbol in self._symbols:
          since = get_timestamp()
          # Query trade history since 30 seconds before.
          since -= ns_before
          fut = self._async_feed_client.async_query_trade(symbol, since)
          if fut is not None:
            self._ioloop.add_future(fut, self.on_trade)
    finally:
      self._notify_close()

  def start(self, ioloop, delay_time=None):
    assert not self.stopped

    self._state = self.State.RUNNING
    self._ioloop = ioloop
    if delay_time is None:
      ioloop.add_callback(self._run_loop)
    else:
      ioloop.call_later(delay_time, self._run_loop)

  def stop(self, trigger_close_callback=True):
    self._async_feed_client.stop()
    super().stop()
